use std::{collections::BinaryHeap, fs::File, iter::zip, path::PathBuf, sync::Arc};

use anyhow::{Error, Result};
use serde::{Deserialize, Serialize};

use crate::{
    content::{ContentCollection, ContentStore},
    train::TrainableEngine,
    SearchConfig, SearchResultItem,
};

use super::{VecEngine, VecEngineBuilder};

pub struct SimpleEngine {
    indexs: Vec<IndexItem>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct IndexItem {
    index: usize,
    vec: Vec<f32>,
}

pub struct InitConfig {
    pub path: PathBuf,
}

impl Default for InitConfig {
    fn default() -> Self {
        Self { path: ".".into() }
    }
}
pub struct TrainConfig {}

impl SimpleEngine {
    fn dis(a: &[f32], b: &[f32]) -> anyhow::Result<f32> {
        if a.len() != b.len() {
            let msg = format!(
                "Vectors should have same length. a has length of {}, while b lengthes {}.",
                a.len(),
                b.len()
            );
            return Err(Error::msg(Arc::new(msg)));
        }
        let dis = zip(a, b).map(|(i, j)| (i - j).powi(2)).sum();
        Ok(dis)
    }
}

impl VecEngineBuilder for SimpleEngine {
    type Engine = Self;

    type InitConfig = InitConfig;

    fn init(config: &Self::InitConfig) -> anyhow::Result<Self>
    where
        Self: Sized,
    {
        let path = &config.path;
        let rdr = File::open(path)?;
        let indexs: Vec<IndexItem> = serde_json::from_reader(rdr)?;
        Ok(Self { indexs })
    }

    fn build(self) -> anyhow::Result<Self::Engine> {
        Ok(self)
    }
}

impl VecEngine for SimpleEngine {
    async fn search_for_ids(
        &self,
        inputs: &[f32],
        contents: &ContentStore,
        config: &SearchConfig,
    ) -> anyhow::Result<Vec<SearchResultItem>> {
        let mut selected = BinaryHeap::new();
        let mut min_dis = f32::MAX;
        let mut max_dis = 0f32;
        let max_result_length = config.max_result_length.unwrap_or(3);
        for inp in &self.indexs {
            let dis = Self::dis(inputs, inp.vec.as_slice())?;
            if selected.len() < max_result_length {
                if dis < min_dis {
                    min_dis = dis;
                }
                if dis > max_dis {
                    max_dis = dis;
                }
                let item = SearchResultItem {
                    id: inp.index,
                    distance: dis,
                };
                selected.push(item);
            } else if dis < max_dis {
                selected.pop();
                let item = SearchResultItem {
                    id: inp.index,
                    distance: dis,
                };
                selected.push(item);
                if let Some(t) = selected.peek() {
                    max_dis = t.distance;
                }
            }
        }
        // let mut result = selected
        //     .into_iter()
        //     .map(|content| content.id)
        //     .collect::<Vec<_>>();
        let mut result = selected.into_vec();
        result.reverse();
        Ok(result)
    }
}

impl TrainableEngine for SimpleEngine {
    type TrainConfig = InitConfig;

    fn blank(_config: &Self::TrainConfig) -> anyhow::Result<Self>
    where
        Self: Sized,
    {
        Ok(Self { indexs: vec![] })
    }

    fn train(&mut self, material: crate::content::ContentStore) -> anyhow::Result<()> {
        let indexs = (*material.get_inner())
            .contents
            .iter()
            .map(|content| IndexItem {
                index: content.inner.id,
                vec: content.inner.vec.clone(),
            })
            .collect::<Vec<_>>();
        self.indexs = indexs;
        Ok(())
    }

    fn save(&self, config: &Self::TrainConfig) -> anyhow::Result<()> {
        let path = &config.path;
        let wtr = File::create(path)?;
        serde_json::to_writer(wtr, &self.indexs)?;
        Ok(())
    }
}
